import os
import re
import json
import time
import logging
from typing import List
from collections import defaultdict

from fastapi import FastAPI
from pydantic import BaseModel
from openai import OpenAI
import uvicorn

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# ============================================================================
# Configuration Management
# ============================================================================
API_KEY = os.environ.get("ARK_API_KEY", "sk-U4BpSvt3NnkGhZpYtZtjmOjWR0YKSkScXVLBk6QIYHaGgaQC")
BASE_URL = os.environ.get("ARK_API_BASE_URL", "http://ai.docker.tcl.com/huizhou/inference/usr-gqb23upd/inf-7qgwnmfo/v1")  # 例如 https://api.openai.com/v1
MODEL_NAME = os.environ.get("ARK_MODEL_NAME", "")

def validate_configuration():
    """Validate that required configuration values are set"""
    errors = []
    
    if not API_KEY or API_KEY == "sk-U4BpSvt3NnkGhZpYtZtjmOjWR0YKSkScXVLBk6QIYHaGgaQC":
        errors.append("ARK_API_KEY is not properly configured")
    
    if not BASE_URL or BASE_URL == "http://ai.docker.tcl.com/huizhou/inference/usr-gqb23upd/inf-7qgwnmfo/v1":
        errors.append("ARK_API_BASE_URL is not properly configured")
    
    return errors

client = OpenAI(
    api_key=API_KEY,
    base_url=BASE_URL
)

# ============================================================================
# Agent Memory Management
# ============================================================================
class Agent_Memory:
    def __init__(self, user_identity: str = ''):
        self.user_identity: str = user_identity
        self.query_record: List[str] = []

    def add_query(self, query: str):
        self.query_record.append(query)

    def merge_memory(self):
        if not self.query_record:
            return "用户还没有提问过问题。"
        history_content = "用户已经进行了以下提问：\n"
        for i, query in enumerate(self.query_record):
            history_content += "第{}次提问：{}\n".format(i+1, query)
        
        print(f"history_content: {history_content}")
        return history_content

# In-memory storage for user conversations
user_memories = defaultdict(Agent_Memory)
# Timestamps for cleanup
memory_timestamps = defaultdict(float)

def get_user_memory(user_identity: str) -> Agent_Memory:
    """Get or create memory for a user"""
    if user_identity not in user_memories:
        user_memories[user_identity] = Agent_Memory(user_identity)
        logger.info(f"Created new memory for user: {user_identity}")
    else:
        logger.debug(f"Retrieved existing memory for user: {user_identity}")
    
    memory_timestamps[user_identity] = time.time()
    return user_memories[user_identity]

def cleanup_old_memories(max_age_seconds: int = 3600):
    """Remove memories older than max_age_seconds"""
    current_time = time.time()
    expired_users = [
        user_id for user_id, timestamp in memory_timestamps.items()
        if current_time - timestamp > max_age_seconds
    ]
    
    if expired_users:
        logger.info(f"Cleaning up {len(expired_users)} expired user memories")
    
    for user_id in expired_users:
        if user_id in user_memories:
            del user_memories[user_id]
        if user_id in memory_timestamps:
            del memory_timestamps[user_id]
        
        logger.debug(f"Removed memory for user: {user_id}")

# ============================================================================
# AI Interaction
# ============================================================================
# 角色设定
agent_role_design = '''你是一个熟悉公司经营管理流程的总裁助理，你和你服务的高管已经共事超过5年；
你非常了解经营者们的工作习惯，当他们向你提出了一些关于公司经营状况的问题后，你不仅能快速用关键数据信息回复并且能够猜到他们后续还会有哪些疑问；
现在你要运用你的能力来完成一些特定的工作。
'''

remake_query_prompt_template = '''一位{}的用户还没有想好下面要向你提出什么问题，你需要给他推荐一些他可能感兴趣的提问。
{}

你推荐的问题不能和用户已经提出的问题出现重复，这种重复指的是语义上的重复而不是仅指内容上的重复；你推荐的问题要确保和之前用户提出的问题在同一个经营领域，不能太过跳脱，导致推荐了用户根本不感兴趣的问题。
请按照下面的json格式输出你推荐的提问，记住不要包含任何json以外的内容：
[
  {{
    "qid": "1",
    "content": ""
  }},
  {{
    "qid": "2",
    "content": ""
  }},
  {{
    "qid": "3",
    "content": ""
  }}
]

推荐至少三个提问，但不要超过八个。
'''

def suggest_agent(memory: Agent_Memory):
    user_id = memory.user_identity
    history = memory.merge_memory()
    
    try:
        logger.info(f"Calling AI model for user: {user_id}")
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": agent_role_design},
                {"role": "user", "content": remake_query_prompt_template.format(user_id, history)}
            ],
            extra_body={"enable_think": False},
            stream=False
        )

        content = response.choices[0].message.content
        logger.info(f"Received response from AI model for user: {user_id}")

        # Pre-compile regex patterns for better performance
        think_tag_pattern = re.compile(r'\<think\>[\s\S]*?\<\/think\>')
        leading_json_pattern = re.compile(r'^(\n+)?```json\n')
        trailing_json_pattern = re.compile(r'\n```$')

        # 清理思维链
        content = think_tag_pattern.sub('', content).strip()
        content = leading_json_pattern.sub('', content, count=1)
        content = trailing_json_pattern.sub('', content, count=1)

        try:
            suggest_json = json.loads(content)
            logger.info(f"Successfully parsed JSON response for user: {user_id}")
        except json.JSONDecodeError:
            logger.error(f"Failed to parse JSON response for user: {user_id}")
            suggest_json = [{"qid": "0", "content": "解析失败，原始输出：" + content}]

        return suggest_json
    except Exception as e:
        # Handle API call errors
        logger.error(f"API call failed for user {user_id}: {str(e)}")
        return [{"qid": "0", "content": f"API调用失败: {str(e)}"}]

# ============================================================================
# FastAPI Application
# ============================================================================
app = FastAPI()

from pydantic import field_validator

class QueryRequest(BaseModel):
    user_identity: str = ""
    user_question: str
    
    @field_validator('user_question')
    def question_must_not_be_empty(cls, v):
        if not v or not v.strip():
            raise ValueError('user_question must not be empty')
        return v
    
    @field_validator('user_identity')
    def identity_must_not_be_too_long(cls, v):
        if len(v) > 100:
            raise ValueError('user_identity must be less than 100 characters')
        return v



from fastapi import HTTPException

@app.post("/suggest")
def suggest_api(req: QueryRequest):
    start_time = time.time()  # 添加这行代码
    try:
        logger.info(f"Received request for user: {req.user_identity}")
        logger.debug(f"User question: {req.user_question}")
        
        # Get or create memory for the user
        memory = get_user_memory(req.user_identity)
        # Add the current query to memory
        memory.add_query(req.user_question)

        suggestions = suggest_agent(memory)
        
        # Clean up old memories periodically
        # In a production system, this should be done in a background task
        cleanup_old_memories()
        
        logger.info(f"Returning {len(suggestions)} suggestions for user: {req.user_identity}")
        end_time = time.time()
        execution_time = end_time - start_time
        logger.info(f"suggest_agent function executed in {execution_time:.4f} seconds for question: {req.user_question}")
        return {"recommendations": suggestions}
    except ValueError as e:
        logger.warning(f"Validation error for user {req.user_identity}: {str(e)}")
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Unexpected error for user {req.user_identity}: {str(e)}")
        raise HTTPException(status_code=500, detail="Internal server error")

# ============================================================================
# Main Function (for testing)
# ============================================================================
def test_suggest_agent():
    """Test the suggest_agent function with a sample memory"""
    # Create a test memory
    record = Agent_Memory()
    record.user_identity = "财务经理"

    user_question = "公司今年的现金流情况怎么样？"
    record.add_query(user_question)

    # Call the suggest_agent function
    suggestions = suggest_agent(record)
    
    # Print results
    print("用户提问：", user_question)
    print("推荐的问题：", json.dumps(suggestions, ensure_ascii=False, indent=2))
    
    # Basic validation
    assert isinstance(suggestions, list), "Suggestions should be a list"
    if suggestions:
        first_suggestion = suggestions[0]
        assert "qid" in first_suggestion, "Each suggestion should have a 'qid' field"
        assert "content" in first_suggestion, "Each suggestion should have a 'content' field"
    
    print("Test passed!")


if __name__ == '__main__':
    #test_suggest_agent()
    uvicorn.run(app, host="0.0.0.0", port=8000)